import cv2
import json
from task_allocator.io import parse_usv_file, parse_port_file, parse_task_file
from task_allocator.allocator import assign_tasks
from task_allocator.visualizer import draw_trajectories
from time import time


def main():
    # 文件路径
    usv_file = 'usv_state.txt'
    port_file = 'ports_coordinates.txt'
    task_file = 'task_list.txt'
    map_file = 'binary_map.png'
    out_json = './time_sequence.json'
    out_img = './trajectories_map.png'

    # 读取地图和输入数据
    map_img = cv2.imread(map_file, cv2.IMREAD_GRAYSCALE)
    usvs = parse_usv_file(usv_file)
    ports = parse_port_file(port_file)
    tasks = parse_task_file(task_file)

    # 分配任务
    start_time = time()
    assignments = assign_tasks(usvs, tasks, lambda_weight=0.1)
    end_time = time()
    print("优化时间:", end_time - start_time)

    # 保存时间序列
    json_result = [{
        "usv_id": a["usv_id"],
        "task_id": a["task_id"],
        "time_sequence": a["time_sequence"]
    } for a in assignments]

    with open(out_json, 'w') as f:
        json.dump(json_result, f, indent=2)
    print(f"结果已保存至: {out_json}")

    # 绘图
    draw_trajectories(map_img, usvs, ports, assignments, out_path=out_img, scale=0.1)


if __name__ == "__main__":
    main()
